import pickle, locale

class BudgetCategory():
    """
    Return the children of a root URL path
    yearTotals returns the total spending amounts for each subtree
        (yearTotals is currently hard coded and only works for a url of level 2:
            url =  "Department_of_Energy/National_Nuclear_Security_Administration" 
        )
    
    Example:
    urlPath = "Department_of_Defense--Military/Operation_and_Maintenance"
    categories = ["Department_of_Defense--Military", "Operation_and_Maintenance"]
    children = ["Aircraft Procurement, Air Force", "Aircraft Procurement, Army", "Aircraft Procurement, Navy", ..]
    yearTotals = "{2000: 599700, 2001: 6084000}"

    """

    def __init__(self, urlPath, startYear=2000, endYear=2009):    
        self.urlPath = urlPath
        self.startYear = startYear
        self.endYear = endYear

      
    def categories(self):
        path = self.urlPath.replace("_", " ")
        return path.split("/")
    
    def children(self):
        pkl_file = open('outlays.pkl', 'rb')   #load Pickled version of budget data
        data = pickle.load(pkl_file)        
        pkl_file.close()

        for i, category in enumerate(self.categories()):
  
            if i < len(self.categories()):
  
                if data.has_key(category):
                    data = data[category]
                else:
                    print("Missing key: ", category)

        return data 
        
    def yearTotals(self): 
        """
        Currently hard coded and only works for a url of level 2:
            url =  "Department_of_Energy/National_Nuclear_Security_Administration" 
        """

        years = {} #Return years and their amounts: {2000: 134000, 2001: 140000}
        
        #Initialize Years to 0
        for i in range(self.startYear,self.endYear + 1):
            years[i] = 0
        
        #Loop through 2 levels of children (hard coded) and add year totals for each subcategory
        for category in self.children():
            
            for subcategory in self.children()[category]:
                
                for fiscalYear, amount in self.children()[category][subcategory].iteritems():
                    if fiscalYear >= self.startYear and fiscalYear <= self.endYear:
                        years[fiscalYear] = years[fiscalYear] + int(amount)        

        return years


    def data(self):
        """
            Returns nested dictionary of all data
        """

        pkl_file = open('outlays.pkl', 'rb')   #load Pickled version of budget data
        data = pickle.load(pkl_file)        
        pkl_file.close()

        return data